import json
import asyncio
import traceback
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import random
from datetime import datetime, timedelta

router = APIRouter(prefix="/chat")

# 活跃的WebSocket连接
active_connections = []

# 模拟聊天历史数据（实际项目中应该从数据库获取）
chat_history = [
    {
        "id": "chat1",
        "title": "系统监控相关问题",
        "date": (datetime.now() - timedelta(days=1)).isoformat(),
        "messages": [
            {"role": "user", "content": "如何查看系统监控信息？"},
            {"role": "assistant", "content": "您可以通过仪表盘查看系统监控信息，包括CPU使用率、内存占用、磁盘I/O等。导航到'监控'页面，您将看到这些信息的图表展示。"}
        ]
    },
    {
        "id": "chat2",
        "title": "性能优化讨论",
        "date": datetime.now().isoformat(),
        "messages": [
            {"role": "user", "content": "怎样优化系统性能？"},
            {"role": "assistant", "content": "优化系统性能可以从多个方面入手：\n1. 关闭不必要的服务和进程\n2. 定期清理缓存和临时文件\n3. 优化数据库查询\n4. 检查并解决资源泄漏问题\n\n您可以在'进程'页面查看当前运行的进程，并关闭不必要的进程。"}
        ]
    }
]

@router.websocket("/ws")
async def websocket_chat(websocket: WebSocket):
    """
    WebSocket处理AI聊天的连接
    """
    await websocket.accept()
    print("AI聊天WebSocket连接已接受")
    active_connections.append(websocket)
    
    try:
        # 发送连接成功消息（仅用于日志记录，不显示在界面上）
        await websocket.send_text(json.dumps({
            "type": "connection_success",
            "content": "WebSocket连接已建立"
        }))
        
        while True:
            # 接收用户消息
            data = await websocket.receive_text()
            print(f"收到聊天消息: {data}")
            
            try:
                # 解析JSON数据
                message = json.loads(data)
                message_type = message.get("type", "")
                
                print(f"处理消息类型: {message_type}")
                
                # 根据消息类型处理
                if message_type == "message":
                    # 处理聊天消息
                    user_input = message.get("content", "")
                    conversation_id = message.get("conversation_id")
                    print(f"用户输入: {user_input}, 会话ID: {conversation_id}")
                    
                    # 模拟AI处理延迟
                    await asyncio.sleep(0.5)
                    
                    # 生成模拟回复
                    response = generate_mock_response(user_input)
                    
                    # 模拟流式响应
                    await stream_response(websocket, response)
                    
                elif message_type == "get_history":
                    # 处理获取历史对话请求
                    conversation_id = message.get("conversation_id")
                    
                    try:
                        if conversation_id:
                            # 获取特定对话的历史
                            conversation = next((chat for chat in chat_history if chat["id"] == conversation_id), None)
                            if conversation:
                                await websocket.send_text(json.dumps({
                                    "type": "history_data",
                                    "content": conversation
                                }))
                            else:
                                # 对于找不到对话的情况，只记录日志，不发送错误消息
                                print(f"未找到ID为{conversation_id}的对话")
                        else:
                            # 获取所有对话的列表（不包含消息内容）
                            history_list = [{
                                "id": chat["id"],
                                "title": chat["title"],
                                "date": chat["date"]
                            } for chat in chat_history]
                            
                            await websocket.send_text(json.dumps({
                                "type": "history_list",
                                "content": history_list
                            }))
                    except Exception as e:
                        # 处理历史数据相关错误，只记录日志，不发送错误消息
                        print(f"处理历史数据时出错: {str(e)}")
                        print(traceback.format_exc())
                
                else:
                    # 对于其他未知类型的消息，只记录日志，不发送错误消息
                    print(f"未知的消息类型: {message_type}")
                
            except json.JSONDecodeError:
                print("JSON解析错误")
                # JSON解析错误，只记录日志，不发送错误消息
            except Exception as e:
                print(f"处理消息时出错: {str(e)}")
                print(traceback.format_exc())
                
                # 只有对聊天消息类型的错误才发送错误响应
                if message_type == "message":
                    await websocket.send_text(json.dumps({
                        "type": "error",
                        "content": f"处理消息时出错: {str(e)}"
                    }))
    except WebSocketDisconnect:
        print("AI聊天WebSocket连接已断开")
    except Exception as e:
        print(f"WebSocket错误: {str(e)}")
        print(traceback.format_exc())
    finally:
        # 移除连接
        if websocket in active_connections:
            active_connections.remove(websocket)
        print("AI聊天WebSocket连接已关闭")

async def stream_response(websocket, response):
    """
    模拟流式响应
    """
    words = response.split()
    current_text = ""
    
    for word in words:
        # 添加当前单词和空格
        current_text += word + " "
        
        # 发送当前累积文本
        await websocket.send_text(json.dumps({
            "type": "stream",
            "content": current_text
        }))
        
        # 模拟打字延迟
        await asyncio.sleep(random.uniform(0.05, 0.2))
    
    # 发送结束信号
    await websocket.send_text(json.dumps({
        "type": "end"
    }))

def generate_mock_response(user_input):
    """
    生成模拟AI回复
    """
    responses = [
        f"您询问的是\"{user_input}\"。这是一个很好的问题！系统监控通常包括CPU使用率、内存占用、磁盘I/O等多个方面。您可以通过仪表盘查看这些数据。",
        f"关于\"{user_input}\"，我建议您检查系统日志文件，它们通常位于/var/log目录下。您也可以使用命令行工具如top、htop来实时监控系统资源。",
        f"对于\"{user_input}\"，最佳实践是建立一个完整的监控体系，包括预警机制和历史数据分析。这样可以帮助您及时发现并解决潜在问题。",
        f"您想了解的\"{user_input}\"涉及到系统安全的重要方面。定期更新系统和应用是保持系统安全的基础措施之一。",
        f"针对\"{user_input}\"，可以考虑使用专业的监控工具，如Prometheus + Grafana组合，它们提供强大的数据收集和可视化能力。"
    ]
    
    # 随机选择一个回复
    return random.choice(responses) 